#!/usr/local/bin/python
# -*- encoding:utf-8 -*-

import numpy as np;
import scipy.sparse as sp
from scipy import linalg
import time
import threading
import os
from scipy.sparse.linalg import eigsh,inv
from multiprocessing import Process, Manager
import sys
import utils
import argparse

'''This file is designed to compute Revised RMLS Algorithms, The target is to minimize following
  
'''

'''This File reads the covariace matrix C_xy, the query synonym words file, 
the document synonym words file'''

config = utils.get_config();

global alpha,beta, converge_threshholad, dx, dy, theta, psi, phi;

alpha = config.getfloat('rmls_plus', 'alpha')
beta = config.getfloat('rmls_plus', 'beta');
converge_threshhold= config.getfloat('rmls_plus', 'converge_threshhold');
dx = config.getint('rmls_plus', 'dx');
dy = config.getint('rmls_plus', 'dy');
d = config.getint('rmls_plus', 'd');
theta = config.getfloat('rmls_plus', 'theta');
psi = config.getfloat('rmls_plus', 'psi');
phi = config.getfloat('rmls_plus', 'phi')
rU_dict = {};
rV_dict = {};

def saveModel(Lx, Ly, modelflag, trainsetflag):
    print 'Saving Model For Batch %s %s' % (modelflag, trainsetflag);
    LxFile, LyFile = utils.spellModelName(modelflag, trainsetflag); 

    utils.save_sparse_csr(LxFile, Lx.tocsr());
    utils.save_sparse_csr(LyFile, Ly.tocsr());
   
def calcC(row, theta):
    if sp.issparse(row):
        t = linalg.norm(row.todense());
    else:
        t = linalg.norm(row);
    if t == 0:
        return 0;
    return theta/t;

def checkConverge(LPx,LPy, Lx, Ly):
    t = np.linalg.norm((Lx-LPx).toarray()) + np.linalg.norm((Ly-LPy).toarray());
    if t < converge_threshhold:
        return (True, t);
    else:
        return (False, t);

def checkConvergeM(LPx, LPy, Lx, Ly):
    t = utils.sparseFroNorm(LPx - Lx) + utils.sparseFroNorm(LPy - Ly);
    if t < converge_threshhold:
        return (True, t);
    else:
        return (False, t);

def computeConvergeM(LP, L):
    return utils.sparseFroNorm(LP -L);

def getRelationGradientLx(u):
    '''Get the \alpha\sum_qL_{xq}'''
    if rU_dict.has_key(u):
        li = rU_dict[u];
        l = len(li);
        if l !=0:
            sum = li[0][1]*Lx[li[0][0]];
        for i in range(1,l):
            sum  = sum + li[i][1]*Lx[li[i][0]];
    else:
        sum = 0;

    return alpha * sum;

def getRelationGradientLy(v):
    '''Get the \beta\sum_qL_{yq}'''
    if rV_dict.has_key(v):
        li = rV_dict[v];

        l = len(li);
        if l != 0:
            sum = li[0][1]*Ly[li[0][0]];
        for i in range(1, l):
            sum = sum + li[i][1]*Ly[li[i][0]];
    else:
        sum = 0;

    return beta * sum;

def updateLxu(slx, Lx, pid, start, end ,debug=False):
    if debug:
        print 'pid: %d, start: %d, end: %d' % (pid, start, end);
    for u in range(start, end):
        if debug:
            start_time = time.time();
        wu = W.getrow(u) + getRelationGradientLx(u);
        lxu = sp.linalg.spsolve(Y, wu.transpose()).transpose();
        #lxu = (wu + getRelationGradientLx(u)).dot(Y);
        #cu = calcC(lxu, 1.0);
        Lx[u] = lxu #* cu;
        if debug:
            print  u, time.time() -start_time;
    slx.append((Lx, start, end));

def updateLyv(sly, Ly, pid, start, end, debug=False):
    if debug:
        print 'pid: %d, start: %d, end: %d' % (pid, start, end);
    for v in range(start, end):
        if debug:
            start_time = time.time();
        nv = N.getrow(v) + getRelationGradientLy(v);
        lyv = sp.linalg.spsolve(X, nv.transpose()).transpose();
        #lyv = (nv + getRelationGradientLy(v)).dot(X);
        #cv = calcC(lyv, 1.0);
        Ly[v] = lyv #* cv;

        if debug:
            print v,time.time() -start_time;
    sly.append((Ly, start, end));

def update(M, sl):
    t = 0;
    for triple in sl:
        R, s, e = triple;
        M[s:e] = R[s:e];
        
    return t;

def RMLS_Plus(C, CT, T, trainsetflag, rpmodelflag):
    con = -1;
    global W, N, Lx, Ly, X, Y;
    K = 25;
    print alpha, beta;
    #Generate Lx and Ly;
    if rpmodelflag == None:
        Lx = sp.rand(dx, d, density=0.1, format='lil', dtype=np.float64);
        Ly = sp.rand(dy, d, density=0.1, format='lil', dtype=np.float64);
    else:
        LxFile, LyFile = utils.spellModelName(rpmodelflag, trainsetflag);
        Lx = utils.load_sparse_csr(LxFile).tolil();
        Ly = utils.load_sparse_csr(LyFile).tolil();

    manager = Manager();
    t = 0
    con = utils.sparseFroNorm(Lx) + utils.sparseFroNorm(Ly);
    Id = sp.identity(d);
    psiId = psi * Id;
    phiId = phi * Id;
    while True:
        t += 1;
        LPx = Lx.tocsr();
        LPy = Ly.tocsr();
        print 'This is %d iteration, last converge is %f' %(t, con);

        print 'Updating Matrix Lx';
        start = time.time();
        W = C.dot(LPy);
        Y = (theta*LPy.transpose().dot(LPy) + psiId);

        ranx = dx/K;
        processListx = [];
        u = 0;
        slx = manager.list();
        for pid in range(K-1):
            processListx.append(Process(target=updateLxu, args=(slx, Lx, pid, u, u+ranx)));
            u = u + ranx;
        processListx.append(Process(target=updateLxu, args=(slx, Lx, K-1, u, dx, False)));

        for pid in range(K):
            processListx[pid].start();
        for pid in range(K):
            processListx[pid].join();
    
        norm =  update(Lx, slx);
        print 'Update Lx cost: ' + str(time.time() -start); 
    
        print 'Updating Matrix Ly';

        start = time.time();
        Lx_csr = Lx.tocsr();
        N = CT.dot(Lx_csr);
        X = (theta*Lx_csr.transpose().dot(Lx_csr) + phiId);
        print X.shape
        rany = dy/K;
        processListy = [];
        v  = 0;
        sly  = manager.list();
        for pid in range(K-1):
            processListy.append(Process(target=updateLyv, args=(sly,Ly,pid, v, v+rany)));
            v = v + rany;
        processListy.append(Process(target=updateLyv, args=(sly, Ly, K-1, v, dy, False)));

        for pid in range(K):
            processListy[pid].start();
        for pid in range(K):
            processListy[pid].join();

        norm +=  update(Ly, sly)    
        print 'Update Ly cost: ' + str(time.time() - start);

        #print LPx, Lx;
        flag, con = checkConvergeM(LPx, LPy, Lx, Ly);
       
        print 'Converge: ', flag, con;
        if flag or t >= T:
            break;
    utils.r2norm(Lx, dx);
    utils.r2norm(Ly, dy);
    return (Lx, Ly);
   
def parse(relational):
    '''this fucntion parse the relational data to a symmetric dict'''
    di = {};

    for line in relational:
        li = line.strip().split('\t');
        i = int(li[0]);
        j = int(li[1]);
        s = float(li[2]);
        if i == -1 or j == -1 or s == 0.0:
            continue;
        if di.has_key(i):
            di[i].append((j,s));
        else:
            di[i] = [];
            di[i].append((j,s));
        if di.has_key(j):
            di[j].append((i,s));
        else:
            di[j] = [];
            di[j].append((i,s));

    return di;

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=' RMLS Plus algorithms for leanring to match from click log and synonyms');
    parser.add_argument('--syn', action='store_true', help='whether to train with synonyms');
    parser.add_argument('-model', '--modelflag', help= 'the runflag of this run');
    parser.add_argument('-train', '--trainsetflag', choices=['five', 'week', 'month'], default='five',  help='specify the testset');
    parser.add_argument('--alpha', type=float);
    parser.add_argument('--beta', type=float);
    parser.add_argument('--d', type=int);
    parser.add_argument('--rpmodelflag', help='this flag can load rmls plus model, and to tune to apply synonyms');
    parser.add_argument('--query', help='query synonyms');
    parser.add_argument('--document', help='document synonyms');
    parser.add_argument('--t', type=int,  help='iteration number');
    args = parser.parse_args();

    if args.syn:
        if args.query:
            query_I = args.query;
        else:
            query_I = 'LabeledData/query_I';

        if args.document:
            document_I = args.document;
        else:
            document_I = 'LabeledData/documents_I';

        relationalU = open(query_I, 'r');
        rU_dict = parse(relationalU);
        relationalU.close();
        relationalV = open(document_I, 'r');
        rV_dict = parse(relationalV);
        relationalV.close();

    rpmodelflag = args.rpmodelflag;

    modelflag = args.modelflag;
    trainsetflag = args.trainsetflag;
    if args.alpha:
        alpha = args.alpha;
    if args.beta:
        beta = args.beta;
    if args.d:
        d = args.d;
    if args.t:
        T = args.t;
    else:
        T = 50;
    WName, NName = utils.spellWNSet(trainsetflag);
    try:
        W = utils.load_sparse_csc(WName); 
        N = W.transpose().tocsc();
        print N.shape;
        start = time.time();
        Lx, Ly = RMLS_Plus(N, W, T, trainsetflag, rpmodelflag); 
        print 'training model cost', time.time() - start;
        saveModel(Lx, Ly, modelflag, trainsetflag);
    except KeyboardInterrupt:
        saveModel(Lx, Ly, modelflag, trainsetflag);
        sys.exit(0);
